#!/bin/bash

#
# tcpdns -- by Dario Berzano <dario.berzano@gmail.com>
#
# Part of the lfscripts suite -- http://lfscripts.googecode.com/
#
# Uses ssh and socat to create a local DNS server listening on 53/udp that
# forwards requests to another remote DNS server, through a ssh tunnel (hence,
# encrypted), transforming UDP to TCP and vice-versa
#

#
# Global variables
#

# The variables UDP_DNS, SSH_STR and TESTHOST must be defined in the config

# Program's variables
export TCP_REMOTE_PORT
export TCP_LOCAL_PORT
export SSH_LN="/tmp/ssh-socat-tcpdns"
export SOCAT_LN="/tmp/socat-tcpdns"

# Colors
export YL='\033[1;33m'
export GR='\033[1;32m'
export MG='\033[1;35m'
export RD='\033[1;31m'
export ZZ='\033[m'
export OK="${GR}OK${ZZ}"
export ER="${RD}FAILED${ZZ}"

# Kill previous server
function KillServer() {

  # Kill ssh session
  echo -ne "${YL}Killing ssh session:${ZZ}"
  local P=$(ps aux|grep -v grep|grep $SSH_LN|cut -b10-14)
  [ "$P" != "" ] && kill -9 $P
  rm -f $SSH_LN
  echo -e " $OK"

  # Kill socat serves
  echo -e "${YL}Killing remote server: ${MG}invoking ssh...${ZZ}"
  ssh $SSH_STR \
    kill -9 '$(ps aux|grep -v grep|grep UDP:'$UDP_DNS':53|cut -b10-14)' \
    > /dev/null 2>&1
  local R=$?
  echo -ne "${YL}Killing remote server: ${ZZ}"
  [ $R != 255 ] && echo -e "$OK" || echo -e "$ER"
  return $R
}

# Kill previous client
function KillClient() {
  echo -ne "${YL}Killing local client:${ZZ}"
  local P=$(ps aux|grep -v grep|grep $SOCAT_LN|cut -b10-14)
  [ "$P" != "" ] && sudo kill -9 $P
  rm -f $SOCAT_LN
  echo -e " $OK"
}

# Server
function Server() {
  echo -e "${YL}Running tunnel and server: ${MG}invoking ssh...${ZZ}"
  ln -nfs `which ssh` $SSH_LN
  $SSH_LN -f $SSH_STR \
    -L$TCP_LOCAL_PORT:localhost:$TCP_REMOTE_PORT \
    socat -v tcp4-listen:$TCP_REMOTE_PORT,reuseaddr,fork UDP:$UDP_DNS:53 \
    > /dev/null 2>&1
  local R=$?
  echo -ne "${YL}Running tunnel and server: ${ZZ}"
  [ $R == 0 ] && echo -e "$OK" || echo -e "$ER"
  return $R
}

# Client
function Client() {
  echo -ne "${YL}Running client:${ZZ}"
  ln -nfs `which socat` $SOCAT_LN
  nohup \
    $SOCAT_LN udp4-listen:53,reuseaddr,fork tcp:localhost:$TCP_LOCAL_PORT \
    > /dev/null 2>&1 &
  local R=$?
  [ $R == 0 ] && echo -e " $OK" || echo -e " $ER"
  return $R
}

# Test
function Test() {
  echo -ne "${YL}Testing DNS on ${MG}${TESTHOST}${YL}:${ZZ}"
  host $TESTHOST 127.0.0.1 > /dev/null 2>&1
  local R=$?
  [ $R == 0 ] && echo -e " $OK" || echo -e " $ER"
  return $R
}

# Random port > 10000
function RandomPort() {
  local R=0
  while [ $R -lt 10000 ]; do
    R=$RANDOM
  done
  echo $R
}

# Load configuration
function LoadConf() {
  echo -ne "${YL}Reading configuration:"
  if [ -e "$HOME/.lfscripts/tcpdns.conf" ]; then
    source $HOME/.lfscripts/tcpdns.conf
    echo -e " $OK"
    return 0
  else
    echo -e " $ER"
  fi
  return 1
}

# Main function
function Main() {

  if [ $(whoami) != "root" ]; then
    echo -e "${MG}Becoming root...${ZZ}"
    sudo "$0" "$@"
    exit $?
  fi

  LoadConf || exit 4

  case "$1" in

    -k|--kill)
      KillClient
      KillServer
    ;;

    *)
      TCP_REMOTE_PORT=`RandomPort`
      TCP_LOCAL_PORT=`RandomPort`
      Server || exit 1
      Client || exit 2
      Test
      if [ $? != 0 ]; then
        KillClient
        KillServer
        exit 3
      fi
    ;;

  esac

}

#
# Entry point
#

Main "$@"
